{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9TV7IYeqifSv"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors. [Licensed under the Apache License, Version 2.0](#scrollTo=ByZjmtFgB_Y5)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kZRlD4utdPuX"
      },
      "outputs": [],
      "source": [
        "%install '.package(url: \"https://github.com/tensorflow/swift-models\", .branch(\"tensorflow-0.11\"))' Datasets ImageClassificationModels\n",
        "print(\"\\u{001B}[2J\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tRIJp_4m_Afz"
      },
      "outputs": [],
      "source": [
        "// #@title Licensed under the Apache License, Version 2.0 (the \"License\"); { display-mode: \"form\" }\n",
        "// Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "// you may not use this file except in compliance with the License.\n",
        "// You may obtain a copy of the License at\n",
        "//\n",
        "// https://www.apache.org/licenses/LICENSE-2.0\n",
        "//\n",
        "// Unless required by applicable law or agreed to in writing, software\n",
        "// distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "// See the License for the specific language governing permissions and\n",
        "// limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sI1ZtrdiA4aY"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        " <td><a target=\"_blank\" href=\"https://www.tensorflow.org/swift/tutorials/introducing_x10\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\">TensorFlow.org에서 보기</a></td>\n",
        " <td><a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/swift/blob/master/docs/site/tutorials/introducing_x10.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\">Google Colab에서 실행하기</a></td>\n",
        " <td><a target=\"_blank\" href=\"https://github.com/tensorflow/swift/blob/master/docs/site/tutorials/introducing_x10.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\">View source on GitHub</a></td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8sa42_NblqRE"
      },
      "source": [
        "# X10 소개\n",
        "\n",
        "기본적으로 TensorFlow용 Swift는 eager 디스패치를 사용하여 텐서 작업을 수행합니다. 이를 통해 빠른 반복이 가능하지만, 머신러닝 모델을 훈련하는 데 가장 성능이 좋은 옵션은 아닙니다.\n",
        "\n",
        "[X10 텐서 라이브러리](https://github.com/tensorflow/swift-apis/blob/master/Documentation/X10/API_GUIDE.md)는 텐서 추적 및 [XLA 컴파일러](https://www.tensorflow.org/xla)를 활용하여 TensorFlow용 Swift에 고성능 백엔드를 추가합니다. 이 튜토리얼에서는 X10을 소개하고 GPU 또는 TPU에서 실행되도록 훈련 루프를 업데이트하는 과정을 안내합니다."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W7MpNcIwIIy8"
      },
      "source": [
        "## Eager vs. X10 텐서"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lM9dRji7IIy8"
      },
      "source": [
        "TensorFlow용 Swift의 가속 계산은 텐서 형식을 통해 수행됩니다. 텐서는 다양한 연산에 참여할 수 있으며 머신러닝 모델의 기본 구성 요소입니다.\n",
        "\n",
        "기본적으로 텐서는 eager 실행을 사용하여 연산별로 계산을 수행합니다. 각 텐서에는 연결된 하드웨어와 여기에 사용되는 백엔드를 설명하는 관련 Device가 있습니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OHRTNQJo1TxT"
      },
      "outputs": [],
      "source": [
        "import TensorFlow\n",
        "import Foundation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FCMWR11NIIy-"
      },
      "outputs": [],
      "source": [
        "let eagerTensor1 = Tensor([0.0, 1.0, 2.0])\n",
        "let eagerTensor2 = Tensor([1.5, 2.5, 3.5])\n",
        "let eagerTensorSum = eagerTensor1 + eagerTensor2\n",
        "eagerTensorSum"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1qad9_yMYf6F"
      },
      "outputs": [],
      "source": [
        "eagerTensor1.device"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HrlMNOinIIy_"
      },
      "source": [
        "GPU 지원 인스턴스에서 이 노트북을 실행하는 경우, 위의 기기 설명에 해당 하드웨어가 반영되어 있어야 합니다. eager 런타임은 TPU를 지원하지 않으므로 TPU 중 하나를 가속기로 사용하는 경우, CPU가 하드웨어 대상으로 사용되는 것을 볼 수 있습니다."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eoyLeSQVIIy9"
      },
      "source": [
        "텐서를 생성할 때 대안을 지정하여 기본 eager 모드 기기를 재정의할 수 있습니다. 이것이 X10 백엔드를 사용하여 계산을 수행하도록 선택하는 방법입니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NrRQhQaHaJm9"
      },
      "outputs": [],
      "source": [
        "let x10Tensor1 = Tensor([0.0, 1.0, 2.0], on: Device.defaultXLA)\n",
        "let x10Tensor2 = Tensor([1.5, 2.5, 3.5], on: Device.defaultXLA)\n",
        "let x10TensorSum = x10Tensor1 + x10Tensor2\n",
        "x10TensorSum"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VbqeudQCaqwv"
      },
      "outputs": [],
      "source": [
        "x10Tensor1.device"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mIbIOW0HIIzA"
      },
      "source": [
        "GPU 지원 인스턴스에서 실행하는 경우, X10 텐서의 기기에 해당 가속기가 나열되어 있어야 합니다. eager 실행과 달리 TPU 지원 인스턴스에서 실행하는 경우, 이제 계산에서 해당 기기를 사용하고 있음을 확인할 수 있습니다. X10은 TensorFlow용 Swift에서 TPU를 활용하는 방법입니다.\n",
        "\n",
        "기본 eager 및 X10 기기는 시스템의 첫 번째 가속기를 사용하려고 시도합니다. GPU가 연결되어 있는 경우, 사용 가능한 첫 번째 GPU를 사용합니다. TPU가 있는 경우, X10은 기본적으로 첫 번째 TPU 코어를 사용합니다. 가속기가 없거나 지원되지 않으면 기본 기기가 CPU로 대체됩니다.\n",
        "\n",
        "기본 eager 및 XLA 기기 외에도 특정 하드웨어 및 백엔드 대상을 Device에 제공할 수 있습니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "De59VwJ35SvG"
      },
      "outputs": [],
      "source": [
        "// let tpu1 = Device(kind: .TPU, ordinal: 1, backend: .XLA)\n",
        "// let tpuTensor1 = Tensor([0.0, 1.0, 2.0], on: tpu1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rU0WY_sJodio"
      },
      "source": [
        "## eager-mode 모델 훈련하기\n",
        "\n",
        "기본 eager 실행 모드를 사용하여 모델을 설정하고 훈련하는 방법을 살펴보겠습니다. 이 예제에서는 [swift-models 리포지토리](https://github.com/tensorflow/swift-models)의 간단한 LeNet-5 모델과 MNIST 숫자 분류 필기 데이터세트를 사용합니다.\n",
        "\n",
        "먼저 MNIST 데이터세트를 설정하고 다운로드합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kqXILiXhq-iM"
      },
      "outputs": [],
      "source": [
        "import Datasets\n",
        "\n",
        "let epochCount = 5\n",
        "let batchSize = 128\n",
        "let dataset = MNIST(batchSize: batchSize)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1pMewsl0VgnJ"
      },
      "source": [
        "다음으로 모델과 옵티마이저를 구성합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fEAEyUExXT3I"
      },
      "outputs": [],
      "source": [
        "import ImageClassificationModels\n",
        "\n",
        "var eagerModel = LeNet()\n",
        "var eagerOptimizer = SGD(for: eagerModel, learningRate: 0.1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cwNDwzS2QgS1"
      },
      "source": [
        "이제 기본 진행 상황 추적 및 보고를 구현합니다. 모든 중간 통계는 훈련이 실행되는 같은 기기에 텐서로 기록되고, `scalarized()`는 보고 중에만 불러냅니다. 이는 지연 텐서의 불필요한 구체화를 방지하기 때문에 나중에 X10을 사용할 때 특히 중요합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eNzOdly3QY_P"
      },
      "outputs": [],
      "source": [
        "struct Statistics {\n",
        "    var correctGuessCount = Tensor<Int32>(0, on: Device.default)\n",
        "    var totalGuessCount = Tensor<Int32>(0, on: Device.default)\n",
        "    var totalLoss = Tensor<Float>(0, on: Device.default)\n",
        "    var batches: Int = 0\n",
        "    var accuracy: Float { \n",
        "        Float(correctGuessCount.scalarized()) / Float(totalGuessCount.scalarized()) * 100 \n",
        "    } \n",
        "    var averageLoss: Float { totalLoss.scalarized() / Float(batches) }\n",
        "\n",
        "    init(on device: Device = Device.default) {\n",
        "        correctGuessCount = Tensor<Int32>(0, on: device)\n",
        "        totalGuessCount = Tensor<Int32>(0, on: device)\n",
        "        totalLoss = Tensor<Float>(0, on: device)\n",
        "    }\n",
        "\n",
        "    mutating func update(logits: Tensor<Float>, labels: Tensor<Int32>, loss: Tensor<Float>) {\n",
        "        let correct = logits.argmax(squeezingAxis: 1) .== labels\n",
        "        correctGuessCount += Tensor<Int32>(correct).sum()\n",
        "        totalGuessCount += Int32(labels.shape[0])\n",
        "        totalLoss += loss\n",
        "        batches += 1\n",
        "    }\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w3lmTRCWT5sS"
      },
      "source": [
        "마지막으로, 5개 epoch 동안 훈련 루프를 통해 모델을 실행합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W9bUsiOxVf_v"
      },
      "outputs": [],
      "source": [
        "print(\"Beginning training...\")\n",
        "\n",
        "for (epoch, batches) in dataset.training.prefix(epochCount).enumerated() {\n",
        "    let start = Date()\n",
        "    var trainStats = Statistics()\n",
        "    var testStats = Statistics()\n",
        "    \n",
        "    Context.local.learningPhase = .training\n",
        "    for batch in batches {\n",
        "        let (images, labels) = (batch.data, batch.label)\n",
        "        let 𝛁model = TensorFlow.gradient(at: eagerModel) { eagerModel -> Tensor<Float> in\n",
        "            let ŷ = eagerModel(images)\n",
        "            let loss = softmaxCrossEntropy(logits: ŷ, labels: labels)\n",
        "            trainStats.update(logits: ŷ, labels: labels, loss: loss)\n",
        "            return loss\n",
        "        }\n",
        "        eagerOptimizer.update(&eagerModel, along: 𝛁model)\n",
        "    }\n",
        "\n",
        "    Context.local.learningPhase = .inference\n",
        "    for batch in dataset.validation {\n",
        "        let (images, labels) = (batch.data, batch.label)\n",
        "        let ŷ = eagerModel(images)\n",
        "        let loss = softmaxCrossEntropy(logits: ŷ, labels: labels)\n",
        "        testStats.update(logits: ŷ, labels: labels, loss: loss)\n",
        "    }\n",
        "\n",
        "    print(\n",
        "        \"\"\"\n",
        "        [Epoch \\(epoch)] \\\n",
        "        Training Loss: \\(String(format: \"%.3f\", trainStats.averageLoss)), \\\n",
        "        Training Accuracy: \\(trainStats.correctGuessCount)/\\(trainStats.totalGuessCount) \\\n",
        "        (\\(String(format: \"%.1f\", trainStats.accuracy))%), \\\n",
        "        Test Loss: \\(String(format: \"%.3f\", testStats.averageLoss)), \\\n",
        "        Test Accuracy: \\(testStats.correctGuessCount)/\\(testStats.totalGuessCount) \\\n",
        "        (\\(String(format: \"%.1f\", testStats.accuracy))%) \\\n",
        "        seconds per epoch: \\(String(format: \"%.1f\", Date().timeIntervalSince(start)))\n",
        "        \"\"\")\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7ED0HGZW2gWY"
      },
      "source": [
        "보시다시피 모델은 예상대로 훈련되었으며 검증 세트에 대한 정확성은 매 epoch마다 증가했습니다. 이것이 바로 TensorFlow용 Swift 모델이 정의되고 eager 실행을 사용하여 실행되는 방법입니다. 이제 X10을 활용하기 위해 어떤 수정 사항이 필요한지 살펴보겠습니다."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Te7sNNx9c_am"
      },
      "source": [
        "## X10 모델 훈련하기\n",
        "\n",
        "데이터세트, 모델 및 옵티마이저에는 기본 eager 실행 기기에서 초기화되는 텐서가 포함되어 있습니다. X10으로 작업하려면 이 텐서를 X10 기기로 이동해야 합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MaN7fM-lAe7r"
      },
      "outputs": [],
      "source": [
        "let device = Device.defaultXLA\n",
        "device"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cJfSg0wDAgC7"
      },
      "source": [
        "데이터세트의 경우 훈련 루프에서 배치가 처리되는 시점에서 이를 수행하므로 eager 실행 모델의 데이터세트를 재사용할 수 있습니다.\n",
        "\n",
        "모델 및 옵티마이저의 경우 eager 실행 기기의 내부 텐서로 초기화한 다음 X10 기기로 이동합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jpcOByipc75O"
      },
      "outputs": [],
      "source": [
        "var x10Model = LeNet()\n",
        "x10Model.move(to: device)\n",
        "\n",
        "var x10Optimizer = SGD(for: x10Model, learningRate: 0.1)\n",
        "x10Optimizer = SGD(copying: x10Optimizer, to: device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hQvza3dUXlr0"
      },
      "source": [
        "훈련 루프에 필요한 수정 사항은 몇 개의 특정 지점에서 이루어집니다. 먼저 훈련 데이터 배치를 X10 기기로 옮겨야 합니다. 이것은 각 배치가 검색될 때 `Tensor(copying:to:)`를 통해 수행됩니다.\n",
        "\n",
        "다음 변경 사항은 훈련 루프 중에 추적을 차단할 위치를 나타내는 것입니다. X10은 코드에 필요한 텐서 계산을 추적하고 해당 추적의 최적화된 표시를 JIT 컴파일하여 동작합니다. 훈련 루프의 경우 추적, 컴파일 및 재사용에 이상적인 섹션으로 같은 연산을 수차례 반복합니다.\n",
        "\n",
        "텐서에서 값을 명시적으로 요청하는 코드가 없으면(일반적으로 `.scalars` 또는 `.scalarized()` 호출로 표시됨) X10은 모든 루프 반복을 컴파일하려고 할 것입니다. 이를 방지하고 특정 지점에서 추적을 잘라내기 위해 옵티마이저가 모델 가중치를 업데이트하고 검증 중에 손실 및 정확성을 얻은 후에 명시적인 `LazyTensorBarrier()`를 배치합니다. 이렇게 하면 두 개의 재사용된 추적 즉, 훈련 루프의 각 단계와 검증 중 추론의 각 배치가 생성됩니다.\n",
        "\n",
        "수정 후 훈련 루프는 다음과 같습니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XrZee8n3Y17_"
      },
      "outputs": [],
      "source": [
        "print(\"Beginning training...\")\n",
        "\n",
        "for (epoch, batches) in dataset.training.prefix(epochCount).enumerated() {\n",
        "    let start = Date()\n",
        "    var trainStats = Statistics(on: device)\n",
        "    var testStats = Statistics(on: device)\n",
        "    \n",
        "    Context.local.learningPhase = .training\n",
        "    for batch in batches {\n",
        "        let (eagerImages, eagerLabels) = (batch.data, batch.label)\n",
        "        let images = Tensor(copying: eagerImages, to: device)\n",
        "        let labels = Tensor(copying: eagerLabels, to: device)\n",
        "        let 𝛁model = TensorFlow.gradient(at: x10Model) { x10Model -> Tensor<Float> in\n",
        "            let ŷ = x10Model(images)\n",
        "            let loss = softmaxCrossEntropy(logits: ŷ, labels: labels)\n",
        "            trainStats.update(logits: ŷ, labels: labels, loss: loss)\n",
        "            return loss\n",
        "        }\n",
        "        x10Optimizer.update(&x10Model, along: 𝛁model)\n",
        "        LazyTensorBarrier()\n",
        "    }\n",
        "\n",
        "    Context.local.learningPhase = .inference\n",
        "    for batch in dataset.validation {\n",
        "        let (eagerImages, eagerLabels) = (batch.data, batch.label)\n",
        "        let images = Tensor(copying: eagerImages, to: device)\n",
        "        let labels = Tensor(copying: eagerLabels, to: device)\n",
        "        let ŷ = x10Model(images)\n",
        "        let loss = softmaxCrossEntropy(logits: ŷ, labels: labels)\n",
        "        LazyTensorBarrier()\n",
        "        testStats.update(logits: ŷ, labels: labels, loss: loss)\n",
        "    }\n",
        "\n",
        "    print(\n",
        "        \"\"\"\n",
        "        [Epoch \\(epoch)] \\\n",
        "        Training Loss: \\(String(format: \"%.3f\", trainStats.averageLoss)), \\\n",
        "        Training Accuracy: \\(trainStats.correctGuessCount)/\\(trainStats.totalGuessCount) \\\n",
        "        (\\(String(format: \"%.1f\", trainStats.accuracy))%), \\\n",
        "        Test Loss: \\(String(format: \"%.3f\", testStats.averageLoss)), \\\n",
        "        Test Accuracy: \\(testStats.correctGuessCount)/\\(testStats.totalGuessCount) \\\n",
        "        (\\(String(format: \"%.1f\", testStats.accuracy))%) \\\n",
        "        seconds per epoch: \\(String(format: \"%.1f\", Date().timeIntervalSince(start)))\n",
        "        \"\"\")\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qej_Z6V3mZnG"
      },
      "source": [
        "X10 백엔드를 사용한 모델 훈련은 이전 eager 실행 모델과 같은 방식으로 진행되어야 합니다. 해당 지점에서 고유한 추적의 JIT 컴파일로 인해 첫 번째 배치 전과 첫 번째 epoch의 끝에서 지연이 발생했을 수 있습니다. 가속기를 부착한 상태에서 이것을 실행한다면, 그 시점 이후의 훈련이 eager 모드보다 빠르게 진행되는 것을 확인했을 것입니다.\n",
        "\n",
        "초기에 추적 컴파일 시간 대비 더 빠른 처리량의 균형을 유지하는 어려움이 있지만, 대부분의 머신러닝 모델에서 반복 작업으로 인한 처리량 증가는 컴파일 오버헤드를 상쇄하고도 남습니다. 실제로 일부 훈련 사례에서 X10으로 처리량이 4배 이상 향상되기도 했습니다.\n",
        "\n",
        "앞에서 언급했듯이, X10을 사용하면 TensorFlow용 Swift 모델을 위한 여러 가속기 클래스를 사용할 수 있게 되어 이제 TPU로 작업하는 것이 가능할 뿐만 아니라 수월해집니다."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "introducing_x10.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Swift",
      "name": "swift"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
